from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, TrainerCallback
import evaluate
import numpy as np
from sklearn.model_selection import train_test_split
# 1. 加载 IMDb 数据集（取小样本用于演示）
dataset = load_dataset("imdb")
# train_dataset = dataset["train"].shuffle(seed=42).select(range(800))
# eval_dataset = dataset["test"].shuffle(seed=42).select(range(200))



# # train_dataset = dataset["train"].shuffle(seed=42)
# # eval_dataset = dataset["test"].shuffle(seed=42)
# 从训练集中划分出 20% 作为验证集
# train_dataset = dataset["train"].train_test_split(test_size=0.2, stratify_by_column="label", seed=42)
# train_dataset = train_dataset["train"]
# eval_dataset = train_dataset["test"]



split = dataset["train"].train_test_split(test_size=0.2, seed=42, stratify_by_column="label")
train_dataset = split["train"]
eval_dataset = split["test"]
#
# print(len(train_dataset), len(eval_dataset))


# 2. 加载模型和 tokenizer
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 3. 数据预处理   truncation=True超出模型长度会截断
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256)

# 数据预处理
train_dataset = train_dataset.map(preprocess_function, batched=True)
eval_dataset = eval_dataset.map(preprocess_function, batched=True)

# 4. 加载模型
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# 5. 定义评估指标 (Accuracy)
accuracy_metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return accuracy_metric.compute(predictions=predictions, references=labels)

# 6. 训练参数
training_args = TrainingArguments(
    output_dir="./results", #指定训练输出（模型权重、日志、评估结果等）保存的目录。
    save_strategy="no",                  # 只有当全部训练结束时才会得到一个最后的模型文件（pytorch_model.bin 等）如果训练中途被中断（比如服务器掉线、OOM、手动停止），就无法从中间恢复，只能重新从头训练
    learning_rate=2e-5, #学习率，表示每次更新模型参数的步长。 2e-5 等于 0.00002，这是 Transformer 模型（如 BERT、RoBERTa）常用的较小学习率，避免训练过程发散。
    warmup_steps=1000,
    per_device_train_batch_size=4, #决定训练阶段每个设备一次送多少样本。
    per_device_eval_batch_size=8,#决定验证阶段每个设备一次送多少样本。
    num_train_epochs=8, #训练总共遍历数据集的轮数（epoch = 全部数据训练 1 次）。
    weight_decay=0.01, #权重衰减（L2 正则化系数），用于防止过拟合。 训练时会在更新参数时额外让权重逐渐减小一点。
    # logging_dir="./logs",#指定训练时的日志文件保存路径（目录）。
    # logging_strategy="steps", #作用：控制多久记录一次日志。这里只设置了 "steps"，意思是按训练步数记录。另一个常见选择是 "epoch"（每轮结束记录一次），或者 "no"（不记录日志）。
    logging_steps=5,#每训练 5 个 step，就会记录一次日志（loss、lr 等信息）到 log 文件里，也会在控制台打
    gradient_accumulation_steps=2,  # 每 4 个小批次进行一次梯度更新
    # 🚀 开启混合精度加速（仅 GPU，有 CUDA 才行）
    fp16=True



)

# 7. 定义 Trainer
trainer = Trainer(
    model=model, #模型
    args=training_args,#训练参数
    train_dataset=train_dataset,#训练参数
    eval_dataset=eval_dataset,#验证/评估数据集。
    tokenizer=tokenizer,#分词器工具
    compute_metrics=compute_metrics,#评估指标

)

# 7.5 Callback：兼容旧版的自动评估
class EvalEveryNSteps(TrainerCallback):
    def __init__(self, n_steps=20, eval_dataset=None):
        self.n_steps = n_steps
        self.eval_dataset = eval_dataset
        self.trainer = None  # 保存 Trainer 实例

    def on_train_begin(self, args, state, control, **kwargs):
        # 保存 trainer 对象，兼容旧版 Transformers
        self.trainer = kwargs.get("trainer", None)

    def on_step_end(self, args, state, control, **kwargs):
        if self.trainer and state.global_step > 0 and state.global_step % self.n_steps == 0:
            print(f"\n[Step {state.global_step}] 开始评估...")
            metrics = self.trainer.evaluate(eval_dataset=self.eval_dataset)
            acc = metrics.get("eval_accuracy", None)
            if acc is not None:
                print(f"[评估结果] Accuracy = {acc:.4f}")
            else:
                print("[评估结果]", metrics)

# 添加回调
trainer.add_callback(EvalEveryNSteps(
    n_steps=20,  # 每 20 步评估一次
    eval_dataset=eval_dataset
))

# 8. 训练 & 评估
trainer.train()
eval_results = trainer.evaluate()
print("最终评估结果：", eval_results)

# 9. 保存最终模型
trainer.save_model("./final_model")